package security

import (
	_ "embed"
	"encoding/json"
	"errors"
	"fmt"
	"os"
	"os/exec"
)

//go:embed contrib/html.tpl
var htmlTemplate string

// ErrTrivyNotFound is returned when the trivy binary is not present in PATH.
var ErrTrivyNotFound = errors.New("trivy not found")

// TrivyReport represents a subset of the Trivy JSON report structure
// used for counting vulnerabilities by severity.
type TrivyReport struct {
	Results []struct {
		Vulnerabilities []struct {
			Severity string `json:"Severity"`
		} `json:"Vulnerabilities"`
	} `json:"Results"`
}

// ScanImage runs Trivy against the specified image and outputs both JSON and HTML reports.
// It prints a summary of found vulnerabilities and optionally fails if critical ones exist.
func ScanImage(image string, failOnCritical bool) error {
	if image == "" {
		return fmt.Errorf("image name is required")
	}

	if _, err := exec.LookPath("trivy"); err != nil {
		return ErrTrivyNotFound
	}

	jsonFile := "scan-result.json"
	htmlFile := "scan-report.html"

	// Generate JSON report
	cmd := exec.Command("trivy", "image", "-f", "json", "-o", jsonFile, image)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return err
	}

	// Generate HTML report using embedded template
	tmp, err := os.CreateTemp("", "trivy-html-*.tpl")
	if err != nil {
		return err
	}
	defer os.Remove(tmp.Name())

	if _, err := tmp.WriteString(htmlTemplate); err != nil {
		tmp.Close()
		return err
	}
	if err := tmp.Close(); err != nil {
		return err
	}

	cmd = exec.Command("trivy", "image", "-f", "template", "--template", "@"+tmp.Name(), "-o", htmlFile, image)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	if err := cmd.Run(); err != nil {
		return err
	}

	data, err := os.ReadFile(jsonFile)
	if err != nil {
		return err
	}

	var report TrivyReport
	if err := json.Unmarshal(data, &report); err != nil {
		return err
	}

	counts := map[string]int{}
	total := 0
	for _, res := range report.Results {
		for _, vuln := range res.Vulnerabilities {
			counts[vuln.Severity]++
			total++
		}
	}

	fmt.Printf("Vulnerabilities found: %d", total)
	if total > 0 {
		fmt.Printf(" (CRITICAL: %d, HIGH: %d, MEDIUM: %d, LOW: %d, UNKNOWN: %d)", counts["CRITICAL"], counts["HIGH"], counts["MEDIUM"], counts["LOW"], counts["UNKNOWN"])
	}
	fmt.Println()

	if failOnCritical && counts["CRITICAL"] > 0 {
		return fmt.Errorf("critical vulnerabilities found")
	}
	return nil
}
